﻿/*  Copyright (C) 2013 AmaroK86 (marcidm 'at' hotmail 'dot' com)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.

 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.

 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
using System;
using System.IO;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using System.Linq;
using System.Text;
using Gibbed.IO;
using System.Windows.Forms;
//using AmaroK86.MassEffect3.PCC;
using AmaroK86.MassEffect3.ZlibBlock;
using AmaroK86.MassEffect3.UnrealObjects;

namespace AmaroK86.MassEffect3.Image
{
    public abstract class ImageFile
    {
        public string fileName = "";
        public string archivePath = null;
        public ImageSize imgSize;
        protected Texture2D.ImageList imgList;
        public int dataSize { get { return (imgData  == null) ? 0 : imgData.Length;  } }
        public byte[] imgData  = null;
        public string format = null;
        public float BPP { get; protected set; }

        protected ImageFile(Texture2D tex2D, string archivePath, ImageSize selectedSize = null, bool includeMipMaps = true)
        {
            this.archivePath = archivePath;
            this.imgSize = (selectedSize == null) ? tex2D.imgList.maxImgSize : selectedSize;
            this.format = (string)tex2D.properties["Format"].Value;
            imgList = new Texture2D.ImageList();
            if (selectedSize == null)
                imgList.AddRange(tex2D.imgList.Where(image => image.storageType != ImageInfo.storage.empty));
            else
            {
                imgList.AddRange(tex2D.imgList.Where(image => image.storageType != ImageInfo.storage.empty && image.imgSize <= selectedSize));
            }

            if (!includeMipMaps)
            {
                ImageInfo uniqueImg = imgList[0];
                imgList.Clear();
                imgList.Add(uniqueImg);
            }
        }

        protected ImageFile(string fileName)
        {

        }

        virtual public byte[] ToArray()
        {
            return null;
        }

        virtual public string subtype()
        {
            return null;
        }

        virtual public byte[] resize()
        {
            return imgData;
        }

        virtual public void saveToFile(string fileName)
        {
        }
    }

    public class DDS : ImageFile
    {
        public enum FourCC : uint
        {
            DXT1 = 0x31545844,
            DXT5 = 0x35545844
        }

        [StructLayout(LayoutKind.Sequential, Pack = 1, Size = 128)]
        private unsafe struct _DDSHeader
        {
            public Int32 magic;
            public Int32 Size;
            public Int32 Flags;
            public UInt32 Height;
            public UInt32 Width;
            public Int32 PitchOrLinearSize;
            public Int32 Depth;
            public Int32 MipMapCount;
            public fixed Int32 Reserved1[11];
            // DDS_PIXELFORMAT
            public Int32 pfSize;
            public Int32 pfFlags;
            public Int32 FourCC;
            public Int32 RGBBitCount;
            public Int32 RBitMask;
            public Int32 GBitMask;
            public Int32 BBitMask;
            public Int32 ABitMask;
            // END DDS_PIXELFORMAT
            public Int32 Caps;
            public Int32 Caps2;
            public Int32 Caps3;
            public Int32 Caps4;
            public Int32 Reserved2;
        }

        private _DDSHeader DDSHeader;
        public bool hasMipMaps { get { return (DDSHeader.MipMapCount != 0); } }

        public DDS(Texture2D tex2D, string archivePath, ImageSize selectedSize = null, bool includeMipMaps = true)
            : base(tex2D, archivePath, selectedSize, includeMipMaps)
        {
            DDSHeader.magic = 0x20534444;
            DDSHeader.Size = 0x7C;
            DDSHeader.Flags = 0x081007; // No Mipmaps
            DDSHeader.Caps = 0x1000; // No Mipmaps
            DDSHeader.pfFlags = 0x4;
            DDSHeader.pfSize = 0x20;
            DDSHeader.Width = imgSize.width;
            DDSHeader.Height = imgSize.height;

            if (imgList.Count > 1)
            {
                DDSHeader.Flags = 0x0A1007; // Mipmaps
                DDSHeader.Caps = 0x401008; // Mipmaps
                DDSHeader.MipMapCount = imgList.Count; // Mipmaps
            }

            switch (format)
            {
                case "PF_DXT1":
                    DDSHeader.FourCC = (int)FourCC.DXT1;
                    BPP = 0.5F;
                    break;
                case "PF_DXT5":
                    DDSHeader.FourCC = (int)FourCC.DXT5;
                    BPP = 1;
                    break;
                case "PF_V8U8":
                    DDSHeader.Flags = 0x021007;
                    DDSHeader.pfFlags = 0x80000;
                    DDSHeader.RGBBitCount = 0x10;
                    DDSHeader.RBitMask = 0x0000FF;
                    DDSHeader.GBitMask = 0x00FF00;
                    BPP = 2;
                    break;
                default: throw new FormatException("Invalid DDS format");
            }

            MemoryStream buffer = new MemoryStream();

            if (archivePath != null)
            {
                using (FileStream archiveStream = File.OpenRead(archivePath))
                {
                    foreach (ImageInfo imgInfo in imgList)
                    {
                        switch (imgInfo.storageType)
                        {
                            case ImageInfo.storage.arcCpr:
                                archiveStream.Seek(imgInfo.offset, SeekOrigin.Begin);
                                buffer.WriteBytes(ZBlock.Decompress(archiveStream, imgInfo.cprSize));
                                break;
                            case ImageInfo.storage.arcUnc:
                                archiveStream.Seek(imgInfo.offset, SeekOrigin.Begin);
                                buffer.WriteBytes(archiveStream.ReadBytes(imgInfo.uncSize));
                                break;
                            /*case Texture2D.storage.pccSto:
                                buffer.WriteBytes(imgInfo.pccImage);
                                break;*/
                        }
                    }
                }
            }

            foreach (ImageInfo imgInfo in imgList)
            {
                if (imgInfo.storageType == ImageInfo.storage.pccSto)
                    buffer.WriteBytes(imgInfo.pccImage);
            }

            imgData = buffer.ToArray();

            // convert DDSHeader to byte array
            /*int headSize = Marshal.SizeOf(DDSHeader);
            headData = new byte[headSize];
            IntPtr ptr = Marshal.AllocHGlobal(headSize);
            Marshal.StructureToPtr(DDSHeader, ptr, true);
            Marshal.Copy(ptr, headData, 0, headData.Length);
            Marshal.FreeHGlobal(ptr);*/
        }

        public DDS(string fileName) : base(fileName) // call the base constructor
        {
            if(Path.GetExtension(fileName).ToLowerInvariant() != ".dds")
                throw new FormatException("invalid file");

            int headerSize = Marshal.SizeOf(DDSHeader);
            FileStream fileStream = new FileStream(fileName, FileMode.Open, FileAccess.Read);
            byte[] headData = new byte[headerSize];
            fileStream.Read(headData, 0, headData.Length);
            imgData = new byte[fileStream.Length - headData.Length];
            fileStream.Read(imgData, 0, imgData.Length);
            fileStream.Close();

            // put data inside DDSHeader
            GCHandle handle = GCHandle.Alloc(headData, GCHandleType.Pinned);
            DDSHeader = (_DDSHeader)Marshal.PtrToStructure(handle.AddrOfPinnedObject(), typeof(_DDSHeader));
            handle.Free();

            imgSize = new ImageSize(DDSHeader.Width,DDSHeader.Height);
            format = subtype();
            switch (format)
            {
                case "PF_DXT1": BPP = 0.5F; break;
                case "PF_DXT5": BPP = 1F; break;
                case "PF_V8U8": BPP = 2F; break;
                default: throw new NotSupportedException("invalid dds file format");
            }
        }

        public override string subtype()
        {
            switch (DDSHeader.FourCC)
            {
                case (int)FourCC.DXT1: return "PF_DXT1";
                case (int)FourCC.DXT5: return "PF_DXT5";
                default: return "PF_V8U8";
            }
        }

        public override byte[] ToArray()
        {
            MemoryStream buffer = new MemoryStream();

            int headSize = Marshal.SizeOf(DDSHeader);
            byte[] head = new byte[headSize];
            IntPtr ptr = Marshal.AllocHGlobal(headSize);

            Marshal.StructureToPtr(DDSHeader, ptr, true);
            Marshal.Copy(ptr, head, 0, headSize);
            Marshal.FreeHGlobal(ptr);

            buffer.WriteBytes(head);
            buffer.WriteBytes(imgData);

            return buffer.ToArray();
        }

        public override void saveToFile(string fileName)
        {
            using (FileStream fileStream = File.Create(fileName))
            {
                fileStream.WriteBytes(ToArray());
            }
        }
    }

    public class TGA : ImageFile
    {
        [StructLayout(LayoutKind.Sequential, Pack = 1, Size = 18)]
        public unsafe struct _TGAHeader
        {
            public byte identsize;
            public byte colourmaptype;
            public byte imagetype;
            public Int16 colourmapstart;
            public Int16 colourmaplength;
            public byte colourmapbits;
            public Int16 xstart;
            public Int16 ystart;
            public Int16 width;
            public Int16 height;
            public byte bits;
            public byte descriptor;
        }

        public _TGAHeader TGAHeader;
        public bool RLECompressed
        { 
            get { return (TGAHeader.imagetype & 0x08) == 0x08; }
            private set { if (value) TGAHeader.imagetype |= 0x08; else TGAHeader.imagetype &= unchecked((byte)~0x08); }
        }
        public bool verticalFlipped
        { 
            get { return (TGAHeader.descriptor & 0x20) == 0x20; }
            private set { if (value) TGAHeader.descriptor |= 0x20; else TGAHeader.descriptor &= unchecked((byte)~0x20); }
        }

        public TGA(Texture2D tex2D, string archivePath, ImageSize selectedSize = null)
            : base(tex2D, archivePath, selectedSize, false)
        {
            TGAHeader.imagetype = 0x2; // RGB Image type
            TGAHeader.width = (short)imgSize.width;
            TGAHeader.height = (short)imgSize.height;
            //TGAHeader.colourmapbits = 32;
            TGAHeader.descriptor = 0x00; // 0x20 = flips the image vertically

            byte[] rawData;
            ImageInfo imgSelected = (selectedSize == null) ? imgList[0] : imgList[selectedSize];
            switch (imgSelected.storageType)
            {
                case ImageInfo.storage.arcCpr:
                    using (FileStream archiveStream = File.OpenRead(archivePath))
                    {
                        archiveStream.Seek(imgSelected.offset, SeekOrigin.Begin);
                        rawData = ZBlock.Decompress(archiveStream, imgSelected.cprSize);
                    }
                    break;
                case ImageInfo.storage.arcUnc:
                    using (FileStream archiveStream = File.OpenRead(archivePath))
                    {
                        archiveStream.Seek(imgSelected.offset, SeekOrigin.Begin);
                        rawData = archiveStream.ReadBytes(imgSelected.uncSize);
                    }
                    break;
                case ImageInfo.storage.pccSto:
                    rawData = imgSelected.pccImage;
                    break;
                default:
                    throw new InvalidDataException("invalid selection image");
            }

            TGAHeader.bits = (byte)(rawData.Length / (imgSize.width * imgSize.height) * 8);
            BPP = TGAHeader.bits / 8;

            switch (format)
            {
                case "PF_G8":
                    {
                        //imgData = rawData;
                        imgData = new byte[rawData.Length * 4];

                        TGAHeader.bits = (byte)(imgData.Length / (imgSize.width * imgSize.height) * 8);
                        BPP = TGAHeader.bits / 8;
                        int intBPP = (int)BPP;

                        for (int i = 0; i < imgSize.height; i++)
                        {
                            for (int j = 0; j < imgSize.width; j++)
                            {
                                imgData[(intBPP * imgSize.width * (imgSize.height - 1 - i)) + (j * intBPP)] = rawData[(imgSize.width * i) + j];
                                imgData[(intBPP * imgSize.width * (imgSize.height - 1 - i)) + (j * intBPP) + 1] = rawData[(imgSize.width * i) + j];
                                imgData[(intBPP * imgSize.width * (imgSize.height - 1 - i)) + (j * intBPP) + 2] = rawData[(imgSize.width * i) + j];
                                imgData[(intBPP * imgSize.width * (imgSize.height - 1 - i)) + (j * intBPP) + 3] = 0xFF;
                            }
                        }
                    }
                    break;
                case "PF_A8R8G8B8":
                    {
                        TGAHeader.descriptor |= 0x08; // 0x08 = 8-bit alpha bits
                        imgData = rawData;
                    }
                    break;
                default: throw new FormatException("Invalid TGA format");
            }

            verticalFlipped = true;
            flipVertically();
            // convert TGAHeader to byte array
            /*int headSize = Marshal.SizeOf(TGAHeader);
            headData = new byte[headSize];
            IntPtr ptr = Marshal.AllocHGlobal(headSize);
            Marshal.StructureToPtr(TGAHeader, ptr, true);
            Marshal.Copy(ptr, headData, 0, headData.Length);
            Marshal.FreeHGlobal(ptr);*/
        }

        public TGA(string fileName)
            : base(fileName) // used when importing file
        {
            if (Path.GetExtension(fileName).ToLowerInvariant() != ".tga")
                throw new FormatException("invalid file");

            int headerSize = Marshal.SizeOf(TGAHeader);
            FileStream fileStream = new FileStream(fileName, FileMode.Open, FileAccess.Read);
            byte[] headData = fileStream.ReadBytes(headerSize);
            imgData = fileStream.ReadBytes((int)fileStream.Length - headerSize);

            // saving tga header data in TGAHeader structure
            GCHandle handle = GCHandle.Alloc(headData, GCHandleType.Pinned);
            TGAHeader = (_TGAHeader)Marshal.PtrToStructure(handle.AddrOfPinnedObject(), typeof(_TGAHeader));
            handle.Free();

            imgSize = new ImageSize((uint)TGAHeader.width, (uint)TGAHeader.height);
            BPP = TGAHeader.bits / 8;

            if (RLECompressed)
                RLEDecompress(); // decompress
            else
                imgData = imgData.Take(TGAHeader.width * TGAHeader.height * (int)BPP).ToArray(); // get rid of tga footer

            if (!verticalFlipped)
                flipVertically();

            format = subtype();
        }

        public override string subtype() // analyze data and gives an "opinion", useful when getting data from an external file
        {
            bool isG8 = true;
            for (int i = 0; i < imgData.Length / BPP; i += (int)BPP)
            {
                isG8 = isG8 && (imgData[i] == imgData[i + 1]) && (imgData[i + 1]  == imgData[i + 2]);
            }
            if (isG8)
                return "PF_G8";
            else
                return "PF_A8R8G8B8";
        }

        public override byte[] resize()
        {
            byte[] buffer;
            int intBPP = (int)BPP;

            switch (format)
            {
                case "PF_G8":
                    buffer = new byte[imgData.Length / 4];
                    for (int i = 0; i < imgData.Length / 4; i++)
                    {
                        buffer[i] = imgData[i * 4];
                    }
                    break;
                default:
                    buffer = imgData;
                    break;
            }
            return buffer;
        }

        public override byte[] ToArray()
        {
            int headSize = Marshal.SizeOf(TGAHeader);
            byte[] head = new byte[headSize];
            IntPtr ptr = Marshal.AllocHGlobal(headSize);

            Marshal.StructureToPtr(TGAHeader, ptr, true);
            Marshal.Copy(ptr, head, 0, headSize);
            Marshal.FreeHGlobal(ptr);

            byte[] total = new byte[headSize + dataSize];
            head.CopyTo(total, 0);
            imgData.CopyTo(total, headSize);
            return total;
        }

        public override void saveToFile(string fileName)
        {
            using (FileStream fileStream = File.Create(fileName))
            {
                fileStream.WriteBytes(ToArray());
            }
        }

        private void RLEDecompress()
        {
            if (RLECompressed) //decompressing image
            {
                using (MemoryStream oldImgData = new MemoryStream(imgData), newImgData = new MemoryStream((int)(imgSize.width * imgSize.height * BPP)))
                {
                    byte[] buffer;

                    while (newImgData.Position < (imgSize.width * imgSize.height * BPP))
                    {
                        int count = oldImgData.ReadByte();
                        bool isCompressed = (count & 0x80) == 0x80; //if value > 128
                        if (isCompressed)
                        {
                            count -= 0x7F; //value - 127 = num of repetitions
                            buffer = new byte[(int)BPP];
                            oldImgData.Read(buffer, 0, buffer.Length);
                            for (int j = 0; j < count; j++) // write pixel for each count
                            {
                                newImgData.Write(buffer, 0, buffer.Length);
                            }
                        }
                        else // if it's not compressed copy the values
                        {
                            buffer = new byte[(int)BPP * (count + 1)];
                            oldImgData.Read(buffer, 0, buffer.Length);
                            newImgData.Write(buffer, 0, buffer.Length);
                        }
                    }
                    if (newImgData.Position != newImgData.Length)
                        throw new BadImageFormatException();
                    RLECompressed = false; //remove RLE flag
                    imgData = newImgData.ToArray();
                }
            }
        }

        protected void flipVertically()
        {
            byte[] imgFlip = new byte[imgData.Length];
            int intBPP = (int)BPP;

            for (int i = 0; i < imgSize.height; i++)
                for (int j = 0; j < imgSize.width; j++)
                    for (int k = 0; k < BPP; k++)
                        imgFlip[(imgSize.width * intBPP * (imgSize.height - 1 - i)) + (j * intBPP) + k] = imgData[(imgSize.width * intBPP * i) + (j * intBPP) + k];

            imgData = imgFlip;
            verticalFlipped = !verticalFlipped;
        }
    }
}
